import random
import numpy as np

from beam.util_psll import get_psll_by_phase


class RISSparrowSearchAlgorithm():
    __bit_num = 0       # 比特数
    __beam_num = 0      # 波束数

    def __init__(self, bit_num, beam_num, population_size=50, num_generations=100, alpha=2, beta=1.5, gamma=1.0):
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        self.population_size = population_size
        self.num_generations = num_generations
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.population = None
        self.best_individual = None
        self.best_fitness = float('inf')
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体

    def fitness(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def apply_boundary_conditions(self, sparrow):
        # 应用边界条件，确保新位置在合法范围内
        return np.clip(sparrow, -180, 180)

    def initialize_population(self, phase_mix_init):
        """初始化种群"""
        self.population = [phase_mix_init] * self.population_size

    # def initialize_population(self, phase_mix_init):
    #     # 初始化第一个狼的位置为phase_mix_init
    #     self.population = [np.copy(phase_mix_init)]
    #     # 对于剩下的狼，基于phase_mix_init进行随机变异
    #     for _ in range(1, self.population_size):
    #         # 生成一个与phase_mix_init形状相同的布尔掩码，每个位置有10%的概率为True
    #         mutation_mask = np.random.rand(*phase_mix_init.shape) < 0.20
    #         # 生成随机偏移量
    #         random_offsets = np.random.uniform(-180, 180, phase_mix_init.shape)
    #         # 创建新的狼位置，并应用布尔掩码进行变异
    #         new_sparrow = np.copy(phase_mix_init)
    #         new_sparrow[mutation_mask] += random_offsets[mutation_mask]
    #         # 应用边界条件
    #         new_sparrow = self.apply_boundary_conditions(new_sparrow)
    #         # 将新位置添加到狼群中
    #         self.population.append(new_sparrow)

    def update_position(self, individual, best_individual, worst_individual, iteration):
        r2 = random.random()
        A = 2 * (self.alpha * r2 - self.alpha) * (iteration / self.num_generations)
        C = 2 * r2
        l = (self.alpha - 1) * random.random() + 1
        p = 2 * (1 - iteration / self.num_generations)
        if random.random() < p:
            # 发现者
            if A >= 1:
                rand_individual = self.population[random.randint(0, self.population_size - 1)]
                D = np.linalg.norm(C * rand_individual - individual)
                new_position = rand_individual - A * D
            else:
                D = np.linalg.norm(best_individual - individual)
                new_position = best_individual - A * D
        else:
            # 加入者
            if A >= 1:
                D = np.linalg.norm(best_individual - worst_individual)
                new_position = (best_individual - D) * np.exp(self.gamma * l * iteration / self.num_generations)
            else:
                D = np.linalg.norm(best_individual - individual)
                new_position = best_individual - D * np.exp(self.gamma * l * iteration / self.num_generations)
        # 警戒者
        if random.random() > 0.5:
            Q = random.randint(0, 1)
            K = random.randint(0, 1) * 2 - 1
            new_position = best_individual + K * (180 - (-180))
        return new_position

    def run(self, phase_mix_init, logger):
        logger.info("population_size=%d, num_generations=%d, alpha=%.1f, beta=%.1f, gamma=%.1f"
                    % (self.population_size, self.num_generations, self.alpha, self.beta, self.gamma))
        self.best_individual = phase_mix_init
        self.best_fitness = self.fitness(phase_mix_init)
        self.best_fitness_history = []
        self.best_individual_history = []
        # 初始化种群
        self.initialize_population(phase_mix_init)
        # 计算初始适应度
        fitness_values = [self.fitness(individual) for individual in self.population]
        # 找到当前最佳和最差个体
        best_index = np.argmin(fitness_values)
        self.best_individual = self.population[best_index]
        self.best_fitness = fitness_values[best_index]
        worst_index = np.argmax(fitness_values)
        worst_individual = self.population[worst_index]
        # 记录历史最佳适应度
        self.best_fitness_history.append(self.best_fitness)
        self.best_individual_history.append(self.best_individual)
        for generation in range(self.num_generations):
            for i in range(self.population_size):
                # 更新位置
                new_position = self.update_position(self.population[i], self.best_individual, worst_individual,
                                                    generation)
                # 边界处理
                new_position = self.apply_boundary_conditions(new_position)
                # 计算新位置的适应度
                new_fitness = self.fitness(new_position)
                # 更新个体
                if new_fitness < fitness_values[i]:
                    self.population[i] = new_position
                    fitness_values[i] = new_fitness
            # 更新最佳和最差个体
            best_index = np.argmin(fitness_values)
            self.best_individual = self.population[best_index]
            self.best_fitness = fitness_values[best_index]
            worst_index = np.argmax(fitness_values)
            worst_individual = self.population[worst_index]
            # 记录历史最佳适应度
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)
            logger.info("generation=%d: self.best_fitness=%f" % (generation, self.best_fitness))
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history
